from core.function import Function
from utils.common import pair

from utils.functions_conv import im2col_array, col2im
import utils

# 执行平均池化操作的函数类
class AveragePooling(Function):
    def __init__(self, kernel_size, stride=1, pad=0):
        """
        平均池化操作的函数类。

        Args:
            kernel_size (int or tuple): 池化窗口的大小，可以是单个整数或包含两个整数的元组。
            stride (int or tuple, optional): 池化操作的步幅，可以是单个整数或包含两个整数的元组。默认为1。
            pad (int or tuple, optional): 池化操作的填充，可以是单个整数或包含两个整数的元组。默认为0。
        """
        super().__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.pad = pad
        self.input_shape = None

    def forward(self, x):
        """
        前向传播，执行平均池化操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 平均池化操作的输出变量。
        """
        self.input_shape = x.shape
        col = im2col_array(x, self.kernel_size, self.stride, self.pad,
                           to_matrix=False)
        y = col.mean(axis=(2, 3))
        return y

    def backward(self, gy):
        """
        反向传播，计算输入变量的梯度。

        Args:
            gy (Variable): 输出变量的梯度。

        Returns:
            Variable: 输入变量的梯度。
        """

        # TODO(Koki): This is a simple implementation
        N, C, OH, OW = gy.shape
        KW, KH = pair(self.kernel_size)
        gy /= (KW * KH)
        gcol = utils.functions_collect.broadcast_to(gy.reshape(-1), (KH, KW, N * C * OH * OW))
        gcol = gcol.reshape(KH, KW, N, C, OH, OW).transpose(2, 3, 0, 1, 4, 5)
        gx = col2im(gcol, self.input_shape, self.kernel_size, self.stride,
                    self.pad, to_matrix=False)
        return gx

    def average_pooling(self, x):
        """
        对输入变量应用平均池化操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 平均池化操作的输出变量。
        """
        return self(x)
